/**
Copyright (c) 2011, Joshua Storck
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
 * The name of Joshua Storck may not be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL JOSHUA STORCK BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

 */

package com.unclestorck.spicfg;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.ServiceLoader;
import java.util.Set;
import java.util.Map.Entry;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Processor;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedOptions;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.MirroredTypeException;
import javax.tools.Diagnostic.Kind;

/**
 * An annotation {@link Processor} that generates a service provider
 * configuration file in META-INF/services/qualifiedServiceName that can be used
 * with {@link ServiceLoader}. This processor will find all classes that are
 * annotated with {@link ServiceProvider} and then group those classes by the
 * {@link ServiceProvider#value()}. For each {@link ServiceProvider#value()},
 * this processor will generate a file named
 * META-INF/services/qualifiedServiceName where qualifiedServiceName is the
 * qualified name of the type specified in {@link ServiceProvider#value()}. The
 * file will contain all of the classes that are annotated with
 * {@link ServiceProvider} for that type. This processor requires an option of
 * 'outputDir' which specifies the output location under which the META-INF is
 * created. This processor will create the META-INF/services directories if they
 * don't already exist.
 * 
 * <p>
 * Here is an example of a class that implements the {@link Processor} service
 * provider interface:
 * 
 * <pre>
 * &#064;ServiceProvider(Processor.class)
 * public class CodeGenerator extends AbstractProcessor {
 *     // Processor methods go here    
 * }
 * </pre>
 * 
 * For the example above, this processor will generate a file named
 * META-INF/services/javax.annotation.processing.Processor which will contain
 * one line with the fully qualified name of the CodeGenerator class.
 * 
 * <p>
 * Unfortunately, the {@link ServiceProviderProcessor} cannot use the
 * {@link ServiceProvider} annotation as it would have to run on itself before
 * it was even compiled (aka chicken and egg problem).
 */
@SupportedOptions( { "outputDir" })
@SupportedSourceVersion(SourceVersion.RELEASE_6)
public class ServiceProviderProcessor extends AbstractProcessor {

    /**
     * Used for obtaining the {@link TypeElement} from the {@link Element} that
     * is returned from the call to
     * {@link RoundEnvironment#getElementsAnnotatedWith(TypeElement)}.
     */
    private static final TypeElementVisitor TYPE_ELEMENT_VISITOR = new TypeElementVisitor();

    /**
     * Used for obtaining the {@link DeclaredType} from the class that is
     * specified in {@link ServiceProvider#value()}.
     */
    private static final DeclaredTypeVisitor DECLARED_TYPE_VISITOR = new DeclaredTypeVisitor();

    /*
     * (non-Javadoc)
     * 
     * @see javax.annotation.processing.AbstractProcessor#process(java.util.Set,
     * javax.annotation.processing.RoundEnvironment)
     */
    @Override
    public boolean process(Set<? extends TypeElement> annotations,
	    RoundEnvironment roundEnv) {
	Map<String, List<TypeElement>> services = new HashMap<String, List<TypeElement>>();
	boolean success = true;
	for (TypeElement annotation : annotations) {
	    if (!ServiceProvider.class.getName().equals(
		    annotation.getQualifiedName().toString()))
		continue;
	    success &= getServices(services, roundEnv
		    .getElementsAnnotatedWith(annotation));
	}
	if (!success || services.isEmpty())
	    return success;
	File servicesDirectory = createServicesDirectory();
	if (servicesDirectory == null)
	    return false;
	for (Entry<String, List<TypeElement>> entry : services.entrySet()) {
	    success &= createServiceFile(servicesDirectory, entry);
	}
	return success;
    }

    /*
     * (non-Javadoc)
     * 
     * @see
     * javax.annotation.processing.AbstractProcessor#getSupportedAnnotationTypes
     * ()
     */
    @Override
    public Set<String> getSupportedAnnotationTypes() {
	return Collections.singleton(ServiceProvider.class.getName());
    }

    /**
     * Creates the META-INF/services directory underneath the directory
     * specified by the outputDir option.
     * 
     * @return the META-INF/services directory or null if there was an error
     */
    private final File createServicesDirectory() {
	String outputDirectoryName = processingEnv.getOptions()
		.get("outputDir");
	if (outputDirectoryName == null) {
	    processingEnv.getMessager()
		    .printMessage(
			    Kind.ERROR,
			    "The outputDir option was not specified for the "
			    + ServiceProvider.class.getName()
				    + " annotation processor");
	    return null;
	}
	File outputDirectory = new File(outputDirectoryName);
	if (!outputDirectory.exists()) {
	    processingEnv.getMessager().printMessage(
		    Kind.ERROR,
		    "The output directory '" + outputDirectoryName
			    + "' for the "
 + ServiceProvider.class.getName()
			    + " annotation processor does not exist");
	    return null;
	}
	File servicesDirectory = new File(outputDirectory, "META-INF");
	servicesDirectory = new File(servicesDirectory, "services");
	if (servicesDirectory.exists())
	    return servicesDirectory;
	else if (!servicesDirectory.mkdirs()) {
	    processingEnv.getMessager().printMessage(
		    Kind.ERROR,
		    "Unable to create services directory '" + servicesDirectory
			    + "'");
	    return null;
	}
	return servicesDirectory;
    }

    /**
     * Populates the specified map with a mapping from a
     * {@link ServiceProvider#value()} to a list of {@link TypeElement}s. This
     * method is mostly for capturing errors due to 'downcasting' from an
     * {@link Element} to a {@link TypeElement} and converting the
     * {@link MirroredTypeException} to an actual {@link TypeElement} because
     * the {@link ServiceProvider#value()} contains a references to a class.
     * 
     * @param services
     *            the populated map
     * @param elements
     *            the elements that are annotated with {@link ServiceProvider}.
     * @return true if it succeeds, false otherwise
     */
    private boolean getServices(Map<String, List<TypeElement>> services,
	    Set<? extends Element> elements) {
	boolean success = true;
	for (Element element : elements) {
	    TypeElement typeElement = element
		    .accept(TYPE_ELEMENT_VISITOR, null);
	    TypeElement serviceType;
	    if (typeElement == null) {
		processingEnv
			.getMessager()
			.printMessage(
				Kind.ERROR,
				"Encountered a "
					+ ServiceProvider.class.getName()
					+ " annotation on a "
					+ element.getKind()
					+ " element when it should have been on a type element",
				element);
		success = false;
	    } else if ((serviceType = getServiceType(typeElement)) != null) {
		List<TypeElement> types = services.get(serviceType
			.getQualifiedName().toString());
		if (types == null) {
		    types = new ArrayList<TypeElement>();
		    services.put(serviceType.getQualifiedName().toString(),
			    types);
		}
		types.add(typeElement);
	    }
	}
	return success;
    }

    /**
     * Returns the {@link TypeElement} for the {@link ServiceProvider#value()}
     * for the {@link ServiceProvider} annotation on the specified
     * {@link TypeElement}. When obtaining an annotation field that is of type
     * {@link Class} in an annotation {@link Processor} a
     * {@link MirroredTypeException} is always thrown. This method converts that
     * {@link MirroredTypeException} to a {@link DeclaredType} and converts that
     * {@link DeclaredType} to a {@link TypeElement}.
     * 
     * @param typeElement
     *            the annotated class
     * @return the {@link ServiceProvider#value()} or null if there is an error
     */
    private final TypeElement getServiceType(TypeElement typeElement) {
	DeclaredType declaredType = null;
	try {
	    typeElement.getAnnotation(ServiceProvider.class).value();
	} catch (MirroredTypeException error) {
	    declaredType = error.getTypeMirror().accept(DECLARED_TYPE_VISITOR,
		    null);
	}
	if (declaredType == null) {
	    processingEnv.getMessager().printMessage(Kind.ERROR, "Service class for the " +
 ServiceProvider.class.getName()
			    + " annotation on "
			    +
		    typeElement.getQualifiedName() + " is not a declared type");
	    return null;
	}
	TypeElement serviceType = declaredType.asElement().accept(
		TYPE_ELEMENT_VISITOR, null);
	if (serviceType == null) {
	    processingEnv.getMessager().printMessage(
		    Kind.ERROR,
		    "Service class for the " + ServiceProvider.class.getName()
			    + " annotation on "
			    + typeElement.getQualifiedName()
			    + " is not a type element");
	    return null;
	}
	return serviceType;
    }

    /**
     * Writes out all of the qualified class names of services to a services
     * file.
     * 
     * @param servicesDirectory
     *            the top level services directory - outputDir/META-INF/services
     * @param entry
     *            the key is the name of the output file and the value is the
     *            list of classes that implement/extend the service.
     * @return true if successful, false otherwise
     */
    private final boolean createServiceFile(File servicesDirectory,
	    Entry<String, List<TypeElement>> entry) {
	File outputFile = new File(servicesDirectory, entry.getKey());
	PrintStream printer;
	try {
	    printer = new PrintStream(new FileOutputStream(outputFile));
	} catch (FileNotFoundException error) {
	    processingEnv.getMessager().printMessage(
		    Kind.ERROR,
		    "Unable to create services file '" + outputFile + "': "
			    + error.getMessage());
	    return false;
	}
	for (TypeElement type : entry.getValue()) {
	    printer.println(type.getQualifiedName());
	}
	printer.close();
	processingEnv.getMessager().printMessage(Kind.NOTE,
		"Generated SPI configuration file '" + outputFile + "'");
	return true;
    }

}
